{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Evaluate LLMs with Hugging Face Lighteval on Amazon SageMaker\n",
    "\n",
    "In this sagemaker example, we are going to learn how to evaluate LLMs using Hugging Face [lighteval](https://github.com/huggingface/lighteval/tree/main).  LightEval is a lightweight LLM evaluation suite that powers [Hugging Face Open LLM Leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard). \n",
    "\n",
    "\n",
    "Evaluating LLMs is crucial for understanding their capabilities and limitations, yet it poses significant challenges due to their complex and opaque nature. LightEval facilitates this evaluation process by enabling LLMs to be assessed on acamedic benchmarks like MMLU or IFEval, providing a structured approach to gauge their performance across diverse tasks.\n",
    "\n",
    "\n",
    "In Detail you will learn how to:\n",
    "1. Setup Development Environment\n",
    "2. Prepare the evaluation configuraiton\n",
    "3. Evaluate Zephyr 7B on TruthfulQA on Amazon SageMaker\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sagemaker --upgrade --quiet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are going to use Sagemaker in a local environment. You need access to an IAM Role with the required permissions for Sagemaker. You can find [here](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) more about it.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "import boto3\n",
    "sess = sagemaker.Session()\n",
    "# sagemaker session bucket -> used for uploading data, models and logs\n",
    "# sagemaker will automatically create this bucket if it not exists\n",
    "sagemaker_session_bucket=None\n",
    "if sagemaker_session_bucket is None and sess is not None:\n",
    "    # set to default bucket if a bucket name is not given\n",
    "    sagemaker_session_bucket = sess.default_bucket()\n",
    "\n",
    "try:\n",
    "    role = sagemaker.get_execution_role()\n",
    "except ValueError:\n",
    "    iam = boto3.client('iam')\n",
    "    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']\n",
    "\n",
    "sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)\n",
    "\n",
    "print(f\"sagemaker role arn: {role}\")\n",
    "print(f\"sagemaker bucket: {sess.default_bucket()}\")\n",
    "print(f\"sagemaker session region: {sess.boto_region_name}\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare the evaluation configuraiton\n",
    "\n",
    "[LightEval](https://github.com/huggingface/lighteval/tree/main) includes script to evaluate LLMs on common benchmarks like MMLU, Truthfulqa, IFEval, and more. It is used to evaluate models on the [Hugging Face Open LLM Leaderboard](https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard). lighteval isy built on top of the great [Eleuther AI Harness](https://github.com/EleutherAI/lm-evaluation-harness) with some additional features and improvements. \n",
    "\n",
    "You can find all available benchmarks [here](https://github.com/huggingface/lighteval/blob/main/examples/tasks/all_tasks.txt). \n",
    "\n",
    "We are going to use Amazon SageMaker Managed Training to evaluate the model. Therefore we will leverage the script available in [lighteval](https://github.com/huggingface/lighteval/blob/main/run_evals_accelerate.py). The Hugging Face DLC is not having lighteval installed. This means need to provide a `requirements.txt` file to install the required dependencies.\n",
    "\n",
    "First lets load the `run_evals_accelerate.py` script and create a `requirements.txt` file with the required dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os \n",
    "import requests as r\n",
    "\n",
    "lighteval_version = \"0.2.0\"\n",
    "\n",
    "# create scripts directory if not exists\n",
    "os.makedirs(\"scripts\", exist_ok=True)\n",
    "\n",
    "# load custom scripts from git\n",
    "raw_github_url = f\"https://raw.githubusercontent.com/huggingface/lighteval/v{lighteval_version}/run_evals_accelerate.py\"\n",
    "res = r.get(raw_github_url)\n",
    "with open(\"scripts/run_evals_accelerate.py\", \"w\") as f:\n",
    "    f.write(res.text)\n",
    "    \n",
    "# write requirements.txt    \n",
    "with open(\"scripts/requirements.txt\", \"w\") as f:\n",
    "    f.write(f\"lighteval=={lighteval_version}\")\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In lighteval, the evaluation is done by running the `run_evals_accelerate.py` script. The script takes a `task` argument which is defined as `suite|task|num_few_shot|{0 or 1 to automatically reduce num_few_shot if prompt is too long}`. Alternatively, you can also provide a path to a txt file with the tasks you want to evaluate the model on, which we are going to do. This makes it easier for you to extend the evaluation to other benchmarks.\n",
    "\n",
    "We are going to evaluate the model on the Truthfulqa benchmark with 0 few-shot examples. [TruthfulQA](https://paperswithcode.com/dataset/truthfulqa) is a benchmark designed to measure whether a language model generates truthful answers to questions, encompassing 817 questions across 38 categories including health, law, finance, and politics​​."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"scripts/tasks.txt\", \"w\") as f:\n",
    "    f.write(f\"lighteval|truthfulqa:mc|0|0\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To evaluate a model on all the benchmarks of the Open LLM Leaderboard you can copy this [file](https://github.com/huggingface/lighteval/blob/v0.2.0/tasks_examples/open_llm_leaderboard_tasks.txt)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Evaluate Zephyr 7B on TruthfulQA on Amazon SageMaker\n",
    "\n",
    "In this example we are going to evaluate the [HuggingFaceH4/zephyr-7b-beta](https://huggingface.co/HuggingFaceH4/zephyr-7b-beta) on the MMLU benchmark, which is part of the Open LLM Leaderboard. \n",
    "\n",
    "In addition to the `task` argument we need to define: \n",
    "* `model_args`: Hugging Face Model ID or path, defined as `pretrained=HuggingFaceH4/zephyr-7b-beta`\n",
    "* `model_dtype`: The model data type, defined as `bfloat16`, `float16` or `float32`\n",
    "* `output_dir`: The directory where the evaluation results will be saved, e.g. `/opt/ml/model` \n",
    "\n",
    "Lightevals can also evaluat peft models or use `chat_templates` you find more about it [here](https://github.com/huggingface/lighteval/blob/v0.2.0/run_evals_accelerate.py). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.huggingface import HuggingFace\n",
    "\n",
    "# hyperparameters, which are passed into the training job\n",
    "hyperparameters = {\n",
    "  'model_args': \"pretrained=HuggingFaceH4/zephyr-7b-beta\", # Hugging Face Model ID\n",
    "  'task': 'tasks.txt',  # 'lighteval|truthfulqa:mc|0|0', \n",
    "  'model_dtype': 'bfloat16', # Torch dtype to load model weights\n",
    "  'output_dir': '/opt/ml/model' # Directory, which sagemaker uploads to s3 after training\n",
    "}\n",
    "\n",
    "# create the Estimator\n",
    "huggingface_estimator = HuggingFace(\n",
    "    entry_point          = 'run_evals_accelerate.py',      # train script\n",
    "    source_dir           = 'scripts',         # directory which includes all the files needed for training\n",
    "    instance_type        = 'ml.g5.4xlarge',   # instances type used for the training job\n",
    "    instance_count       = 1,                 # the number of instances used for training\n",
    "    base_job_name        = \"lighteval\",          # the name of the training job\n",
    "    role                 = role,              # Iam role used in training job to access AWS ressources, e.g. S3\n",
    "    volume_size          = 300,               # the size of the EBS volume in GB\n",
    "    transformers_version = '4.36',            # the transformers version used in the training job\n",
    "    pytorch_version      = '2.1',            # the pytorch_version version used in the training job\n",
    "    py_version           = 'py310',            # the python version used in the training job\n",
    "    hyperparameters      =  hyperparameters,\n",
    "    environment          = { \n",
    "                            \"HUGGINGFACE_HUB_CACHE\": \"/tmp/.cache\",\n",
    "                            # \"HF_TOKEN\": \"REPALCE_WITH_YOUR_TOKEN\" # needed for private models\n",
    "                            }, # set env variable to cache models in /tmp\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now start our evaluation job, with the `.fit()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# starting the train job with our uploaded datasets as input\n",
    "huggingface_estimator.fit()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After the evaluation job is finished, we can download the evaluation results from the S3 bucket. Lighteval will save the results and generations in the `output_dir`. The results are savedas json and include detailed information about each task and the model's performance. The results are available in the `results` key. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "import json\n",
    "import io\n",
    "import os\n",
    "from sagemaker.s3 import S3Downloader\n",
    "\n",
    "\n",
    "# download results from s3\n",
    "results_tar = S3Downloader.read_bytes(huggingface_estimator.model_data)\n",
    "model_id = hyperparameters[\"model_args\"].split(\"=\")[1]\n",
    "result={}\n",
    "\n",
    "# Use tarfile to open the tar content directly from bytes\n",
    "with tarfile.open(fileobj=io.BytesIO(results_tar), mode=\"r:gz\") as tar:\n",
    "    # Iterate over items in tar archive to find your json file by its path\n",
    "    for member in tar.getmembers():\n",
    "        # get path of results based on model id used to evaluate\n",
    "        if os.path.join(\"details\", model_id) in member.name and member.name.endswith('.json'):\n",
    "            # Extract the file content\n",
    "            f = tar.extractfile(member)\n",
    "            if f is not None:\n",
    "                content = f.read()\n",
    "                result = json.loads(content)\n",
    "                break\n",
    "            \n",
    "# print results\n",
    "print(result[\"results\"])\n",
    "# {'lighteval|truthfulqa:mc|0': {'truthfulqa_mc1': 0.40636474908200737, 'truthfulqa_mc1_stderr': 0.017193835812093897, 'truthfulqa_mc2': 0.5747003398184238, 'truthfulqa_mc2_stderr': 0.015742356478301463}}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In our test we achieved a `mc1` score of 40.6% and an `mc2` score of 57.47%. The `mc2` is the score used in the Open LLM Leaderboard. Zephyr 7B achieved a `mc2` score of 57.47% on the TruthfulQA benchmark, which is identical to the score on the Open LLM Leaderboard. \n",
    "The evaluation on Truthfulqa took `999 seconds`. The ml.g5.4xlarge instance we used costs `$2.03 per hour` for on-demand usage. As a result, the total cost for evaluating Zephyr 7B on Truthfulqa was `$0.56`."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
